package com.hliushi.spark.sql

import org.apache.spark.sql.expressions.{Window, WindowSpec}
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
 * descriptions: 窗口函数 topN问题
 *
 * author: Hliushi
 * date: 2021/5/20 14:54
 */
object WindowFun {


  /**
   * 从数据集中得到每个类别收入第一的商品和收入第二的商品
   *
   * @param args
   */
  def main(args: Array[String]): Unit = {
    // 1.创建SparkSession
    val spark: SparkSession = SparkSession.builder()
      .appName("window")
      .master("local[6]")
      .getOrCreate()

    import spark.implicits._

    val source: DataFrame = Seq(
      ("Thin", "Cell phone", 6000),
      ("Norma1", "Tablet", 1500),
      ("Mini", "Tablet", 5500),
      ("Uitra thin", "Cell phone", 5000),
      ("Very thin", "Cell phone", 6000),
      ("Big", "Tablet", 2500),
      ("Bendab1e", "Cell phone", 3000),
      ("Fo1dable", "Cell phone", 3000),
      ("Pro", "Tablet", 4500),
      ("Pro2", "Tablet", 6500)
    ).toDF("product", "category", "revenue")

    import org.apache.spark.sql.functions._

    // 1.定义窗口
    /**
     * WindowSpec : 窗口的描述符, 描述窗口应该是怎么样的
     * dense_rank() over window : 表示一个叫做dense_rank()的函数作用于每一个窗口
     */
    val windowSpec: WindowSpec = Window.partitionBy($"category")
      .orderBy($"revenue".desc)

    // 2.数据处理
    // def dense_rank(): Column
    // def rank(): Column
    // def row_number(): Column
    source.select($"product", $"category", $"revenue", dense_rank() over (windowSpec) as "rank")
      .where($"rank" <= 2)
      .show()

    //  +----------+----------+-------+----+
    //  |   product|  category|revenue|rank|
    //  +----------+----------+-------+----+
    //  |      Thin|Cell phone|   6000|   1|
    //  | Very thin|Cell phone|   6000|   1|
    //  |Uitra thin|Cell phone|   5000|   2|
    //  |      Pro2|    Tablet|   6500|   1|
    //  |      Mini|    Tablet|   5500|   2|
    //  +----------+----------+-------+----+

    // 定义窗口方式2

    source.createOrReplaceTempView("productRevenue")
    val sqlStr =
      """
        |select
        |   tmp.product, tmp.category, tmp.revenue, tmp.rank_rn
        |from (
        |     select
        |        product, category, revenue,
        |        row_number() over(partition by category order by revenue desc) as rank_rn
        |     from
        |        productRevenue
        |     ) as tmp
        |where tmp.rank_rn <= 2
        |""".stripMargin

    spark.sql(sqlStr).show()
    //  +---------+----------+-------+-------+
    //  |  product|  category|revenue|rank_rn|
    //  +---------+----------+-------+-------+
    //  |     Thin|Cell phone|   6000|      1|
    //  |Very thin|Cell phone|   6000|      2|
    //  |     Pro2|    Tablet|   6500|      1|
    //  |     Mini|    Tablet|   5500|      2|
    //  +---------+----------+-------+-------+
  }


}